package com.xiaoze.state;

import org.apache.flink.api.common.functions.RichFlatMapFunction;
import org.apache.flink.api.common.state.ValueState;
import org.apache.flink.api.common.state.ValueStateDescriptor;
import org.apache.flink.api.common.typeinfo.TypeHint;
import org.apache.flink.api.common.typeinfo.TypeInformation;
import org.apache.flink.api.java.tuple.Tuple2;
import org.apache.flink.configuration.Configuration;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.flink.util.Collector;

import java.io.IOException;

/**
 * TestValueState
 *
 * @author Lin Zemin
 * @date 2021-12-30
 */
public class TestValueState {

    public static void main(String[] args) throws Exception {

        StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
        //设置全局并行度，如果不设置，则默认为配置文件中配置的平行度。优先级：代码>配置文件
        env.setParallelism(1);

        env.fromElements(
                        Tuple2.of(1L, 3L),
                        Tuple2.of(2L, 5L),
                        Tuple2.of(1L, 6L),
                        Tuple2.of(2L, 7L),
                        Tuple2.of(1L, 2L),
                        Tuple2.of(2L, 7L),
                        Tuple2.of(1L, 9L),
                        Tuple2.of(1L, 7L),
                        Tuple2.of(1L, 3L),
                        Tuple2.of(1L, 2L))
                .keyBy(key -> key.f0)
                .flatMap(new CountWindowAverage())
                .print();

        env.execute();

    }

    private static class CountWindowAverage extends RichFlatMapFunction<Tuple2<Long, Long>, Tuple2<Long, Long>> {

        private static final long serialVersionUID = 8121064747076974609L;

        private ValueState<Tuple2<Integer, Long>> sum;

        @Override
        public void open(Configuration parameters) {
            ValueStateDescriptor<Tuple2<Integer, Long>> descriptor = new ValueStateDescriptor<>(
                    "average",
                    TypeInformation.of(new TypeHint<Tuple2<Integer, Long>>() {})
            );
            sum = getRuntimeContext().getState(descriptor);
        }

        @Override
        public void flatMap(Tuple2<Long, Long> input, Collector<Tuple2<Long, Long>> collector) throws IOException {

            System.out.println("当前input数据为：" + input);

            Tuple2<Integer, Long> currentSum = sum.value();
            System.out.println("当前currentSum数据为：" + currentSum);

            if(currentSum == null) {
                currentSum = new Tuple2<>(0, 0L);
            }

            currentSum.f0++;
            currentSum.f1 += input.f1;
            sum.update(currentSum);

            System.out.println("更新后currentSum数据为：" + currentSum);

            if(currentSum.f0 >= 2) {
                // currentSum.f0其实就是2
                collector.collect(new Tuple2<>(input.f0, currentSum.f1 * currentSum.f0));
                sum.clear();
                System.out.println("结束后sum数据为：" + sum.value());
            }
            System.out.println();
        }
    }

}
